import gymnasium as gym
from Hunt_Env.Hunt_env.envs.Hunt_world import HunterEnv
from Hunt_Env.Hunt_env.envs.decision.rule_based_decision import rule_based_decision
from gymnasium.wrappers import RecordVideo
from time import sleep

# 渲染环境
env= HunterEnv(render_mode = 'human')

# from line_profiler import LineProfiler    
# lp = LineProfiler()
# lp.enable()
# lp.add_function(env.step)
# lp.add_function(env.reset)
# lp.add_function(env._get_obs)
# lp.add_function(env._get_info)
win_hunter = 0
win_escaper = 0
for i in range(100):
    obs = env.reset()
    terminated = False
    truncated = False
    while not terminated and not truncated:
        hunter_obs = obs['agent_0']
        escaper_obs = obs['agent_1']
        action_hunter = rule_based_decision(hunter_obs, '猎手')
        action_escaper = rule_based_decision(escaper_obs, '逃脱者')
        obs, reward, terminated, truncated, info = env.step({
            'agent_0': action_hunter,
            'agent_1': action_escaper
        })
        # print(env.current_step)
        # env.render()
        if terminated:
            win_hunter += 1
        if truncated:
            win_escaper += 1
        # sleep(0.1)
print(f'猎人胜利次数：{win_hunter}')
print(f'逃脱者胜利次数：{win_escaper}')
# lp.disable()
# lp.print_stats()
